package sample.methodremoveinstruct;

import org.objectweb.asm.*;
import util.FileUtils;

import java.io.IOException;

import static org.objectweb.asm.Opcodes.*;
import static util.StateMachineUtils.SEEN_NOTHING;

public class MethodRemoveZeroAddDemo {

    public static void main(String[] args) throws IOException {
        Class<?> clazz = HelloWorld.class;
        byte[] bytes = invoke(clazz);
        //  output
        FileUtils.writeBytes(Type.getInternalName(clazz) + "RemoveZeroAdd.class", bytes);
        //  replace
        FileUtils.writeBytes(clazz.getResource("/").getPath() + Type.getInternalName(clazz) + ".class", bytes);
    }

    private static byte[] invoke(Class<?> clazz) throws IOException {
        ClassReader classReader = new ClassReader(clazz.getName());
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES);
        MethodRemoveZeroAddVisitor visitor = new MethodRemoveZeroAddVisitor(Opcodes.ASM9, classWriter);
        classReader.accept(visitor, ClassReader.EXPAND_FRAMES);
        return classWriter.toByteArray();
    }

    protected static class MethodRemoveZeroAddVisitor extends ClassVisitor {
        public MethodRemoveZeroAddVisitor(int api, ClassVisitor classVisitor) {
            super(api, classVisitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            if (mv != null && !"<init>".equals(name) && !"<clinit>".equals(name)) {
                boolean isAbstractMethod = (access & ACC_ABSTRACT) != 0;
                boolean isNativeMethod = (access & ACC_NATIVE) != 0;
                if (!isAbstractMethod && !isNativeMethod) {
                    mv = new MethodRemoveAddZeroAdapter(api, mv);
                }
            }
            return mv;
        }

        private class MethodRemoveAddZeroAdapter extends MethodPatternAdapter {
            private static final int SEEN_ICONST_0 = 1;

            public MethodRemoveAddZeroAdapter(int api, MethodVisitor mv) {
                super(api, mv);
            }

            @Override
            public void visitInsn(int opcode) {
                if (state == SEEN_NOTHING) {
                    if (opcode == ICONST_0) {
                        state = SEEN_ICONST_0;
                        return;
                    }
                }
                if (state == SEEN_ICONST_0) {
                    if (opcode == IADD) {
                        state = SEEN_NOTHING;
                        return;
                    }
                    if (opcode == ICONST_0) {
                        mv.visitInsn(ICONST_0);
                        return;
                    }
                }
                super.visitInsn(opcode);
            }

            @Override
            protected void visitInsn() {
                if (state == SEEN_ICONST_0) {
                    mv.visitInsn(ICONST_0);
                }
                state = SEEN_NOTHING;
            }
        }
    }
}
